import mtcnn.detect_face as detect_face
import argparse
import os
from time import sleep
import random
import sys
import tensorflow as tf
import numpy as np
from scipy import misc
import imageio
import cv2

def to_rgb(img):
    w, h = img.shape
    ret = np.empty((w, h, 3), dtype=np.uint8)
    ret[:, :, 0] = ret[:, :, 1] = ret[:, :, 2] = img
    return ret
class ImageClass():
    def __init__(self, name, image_paths):
        self.name = name
        self.image_paths = image_paths

    def __str__(self):
        return self.name + ',' + str(len(self.image_paths)) + 'images'

    def __len__(self):
        return len(self.image_paths)
def get_image_paths(facedir):
    img_path = []
    if os.path.isdir(facedir):
        imgs = os.listdir(facedir)
        img_path = [os.path.join(facedir, img) for img in imgs]
    return img_path
def get_dataset(path):
    dataset = []
    path_exp = os.path.expanduser(path)
    classes = [path for path in os.listdir(path_exp) if os.path.isdir(os.path.join(path_exp, path))]
    classes.sort()
    n_classes = len(classes)
    for i in range(n_classes):
        class_name = classes[i]
        face_dir = os.path.join(path_exp, class_name)
        img_path = get_image_paths(face_dir)
        dataset.append(ImageClass(class_name, img_path))
    return dataset
def store_revision_info(output_dir, arg_string):
    rev_info_filename = os.path.join(output_dir, 'revision_info.txt')
    with open(rev_info_filename, 'w') as f:
        f.write('arguments: %s\n--------------------\n' % arg_string)
        f.write('tensorflow version: %s\n--------------------\n' % tf.__version__)

def parse_arguments(argv):
    parser = argparse.ArgumentParser()

    parser.add_argument('--image_path', type=str, help='Unaligned face images.', default=r'E:\testDir\ml\detect_face.jpg')
    parser.add_argument('--output_dir', type=str, help='Test MTCNN output dir',
                        default=r'E:\testDir\ml\mtcnn_output')
    parser.add_argument('--image_size', type=int, help='Image size (height, width) in pixels.', default=182)
    parser.add_argument('--margin', type=int,
                        help='Margin for the crop around the bounding box (height, width) in pixels.', default=44)
    parser.add_argument('--random_order',
                        help='Shuffles the order of images to enable alignment using multiple processes.',
                        action='store_true')
    parser.add_argument('--gpu_memory_fraction', type=float,
                        help='Upper bound on the amount of GPU memory that will be used by the process.', default=1.0)
    parser.add_argument('--detect_multiple_faces', type=bool,
                        help='Detect and align multiple faces per image.', default=False)

    return parser.parse_args(argv)
def main(args):
    sleep(random.random())
    print("image_path:%s" % args.image_path)
    print("output_dir:%s" % args.output_dir)
    print("image_size:%s" % args.image_size)
    print("margin:%s" % args.margin)
    print("random_order:%s" % args.random_order)
    print("gpu_memory_fraction:%s" % args.gpu_memory_fraction)
    print("detect_multiple_faces:%s" % args.detect_multiple_faces)

    output_dir = os.path.expanduser(args.output_dir)
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # store_revision_info(output_dir, ' '.join(sys.argv))

    print('Creating networks and loading parameters!')

    with tf.Graph().as_default():
        gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=args.gpu_memory_fraction)
        sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options, log_device_placement=False))
        with sess.as_default():
            pnet, rnet, onet = detect_face.create_mtcnn(sess, None)
            saver = tf.train.Saver()
            model_dir = os.path.join(args.output_dir, 'saved_model_for_serving')
            if not os.path.exists(model_dir):
                os.makedirs(model_dir)
            detect_face.save_mtcnn(sess, saver, model_dir)
            log_dir = os.path.join(args.output_dir, 'saved_model_for_log')
            if not os.path.exists(log_dir):
                os.makedirs(log_dir)
            writer = tf.summary.FileWriter(log_dir, sess.graph)
            writer.close()

    minsize = 20  # minimum size of face
    threshold = [0.6, 0.7, 0.7]  # three steps's threshold
    factor = 0.709  # scale factor
    # Add a random key to the filename to allow align ment using multiple processes
    random_key = np.random.randint(0, high=99999)
    bounding_boxes_filename = os.path.join(output_dir, 'bounding_boxes_%05d.txt' % random_key)
    image_path = args.image_path
    filename = os.path.splitext(os.path.split(image_path)[1])[0]
    output_filename = os.path.join(output_dir, filename + '.png')
    try:
        img = imageio.imread(image_path)
    except (IOError, ValueError, IndexError) as e:
        errorMessage = '{}: {}'.format(image_path, e)
        print(errorMessage)
    else:
        if img.ndim < 2:
            print('Unable to align "%s"' % image_path)
        if img.ndim == 2:
            img = to_rgb(img)
        img = img[:, :, 0:3]

        bounding_boxes, key_points = detect_face.detect_face(img, minsize, pnet, rnet, onet, threshold, factor)
        nrof_faces = bounding_boxes.shape[0]
        if nrof_faces > 0:
            det = bounding_boxes[:, 0:4]
            img_size = np.asarray(img.shape)[0:2]
            if nrof_faces > 1:
                bounding_box_size = (det[:, 2] - det[:, 0]) * (det[:, 3] - det[:, 1])
                img_center = img_size / 2
                offsets = np.vstack(
                    [(det[:, 0] + det[:, 2]) / 2 - img_center[1], (det[:, 1] + det[:, 3]) / 2 - img_center[0]])
                offset_dist_squared = np.sum(np.power(offsets, 2.0), 0)
                index = np.argmax(bounding_box_size - offset_dist_squared * 2.0)  # some extra weight on the centering
                det = det[index, :]
            det = np.squeeze(det)
            bb = np.zeros(4, dtype=np.int32)
            print(det[0], det[1], det[2], det[3])
            bb[0] = np.maximum(det[0] - args.margin / 2, 0)
            bb[1] = np.maximum(det[1] - args.margin / 2, 0)
            bb[2] = np.minimum(det[2] + args.margin / 2, img_size[1])
            bb[3] = np.minimum(det[3] + args.margin / 2, img_size[0])
            # print(bb)
            # cropped = img[bb[1]:bb[3], bb[0]:bb[2], :]
            # scaled = misc.imresize(cropped, (args.image_size, args.image_size), interp='bilinear')
            # misc.imsave(output_filename, scaled)
            showimage = cv2.imread(image_path)
            cv2.rectangle(showimage, (int(det[0]),int(det[1])), (int(det[2]),int(det[3])), (0, 255, 0), 2)
            key_points = np.squeeze(key_points)
            print(key_points)
            for i in range(int(len(key_points) / 2)):
                key_point = (int(key_points[i]),int(key_points[i + 4]))
                cv2.circle(showimage, key_point, 1, (0, 255, 0), 4)
            cv2.imwrite(output_filename,showimage)
        else:
            print('Unable to align "%s"' % image_path)

if __name__ == '__main__':
    main(parse_arguments(sys.argv[1:]))